import torch


def train(net, optimizer, device, criterion, train_loader):
    """训练"""
    # 1、设置模型状态为训练状态
    net.train(True)

    # 2、所有梯度置0
    net.zero_grad()

    batch_num = len(train_loader)
    running_loss = 0.0
    for i, data in enumerate(train_loader, start=1):
        # 将输入传入CPU
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        # 3、所有梯度置0
        # 根据pytorch中的backward()函数的计算，当网络参量进行反馈时，梯度是被积累的而不是被替换掉
        # 在每一个batch时不需要将两个batch的梯度混合起来累积，因此这里就需要每个batch设置一遍zero_grad了
        optimizer.zero_grad()

        # 4、网络的前向传播model(inputs)
        outputs = net(inputs)

        # 5、得到损失
        # criterion(outputs, labels)，然后将输出的outputs和原来导入的labels作为loss函数的输入就可以得到损失了
        loss = criterion(outputs, labels)

        # 6、预测该样本属于哪个类别的信息
        # 得到网络的全连接层的输出后，采用torch.max得到模型预测该样本属于哪个类别的信息
        # torch.max()的第一个输入是tensor格式，所以用outputs.data而不是outputs作为输入；第二个参数1是代表dim的意思，也就是取每一行的最大值
        _, preds = torch.max(outputs.data, 1)

        # 7、回传损失
        # 计算得到loss后就要回传损失。这是在训练的时候才会有的操作，测试时候只有forward过程。
        loss.backward()

        # 8、根据这些梯度更新参数
        # 回传损失过程中会计算梯度，然后需要根据这些梯度更新参数，optimizer.step()就是用来更新参数的。optimizer.step()后，你就可以从optimizer.param_groups[0][‘params’]里面看到各个层的梯度和权值信息。
        optimizer.step()

        # 计算误差并显示
        running_loss += loss.item()
        if i % 10 == 0:
            print(preds)
            print('batch:{}/{} loss:{:.3f}'.format(i, batch_num, running_loss / 20))
            running_loss = 0.0